#!/usr/bin/env python
# -*- coding: utf-8; py-indent-offset:4 -*-

# @Time    :2023/04/30 17:43:20
# @Author  :wakeblade (2390245#qq.com) 
# @version :8.1

"""
crypto.py -- 涉及到密码生成及加密解密的方法 
"""

import random
import secrets
import string
import hmac
import hashlib
import base64
from Crypto.Cipher import AES
from Crypto import Random
from Crypto.Util.Padding import pad, unpad

alphanum = string.ascii_letters + string.digits

def simple(len: int) -> str:
    """生成固定长度随机密码

    Args:
        len (int): 密码长度

    Returns:
        str: 随机密码
    """
    return ''.join(secrets.choice(alphanum) for i in range(len))

def strong(len: int, min_uppers:int = 1, min_lowers:int = 1, min_digits:int = 3) -> str:
    """生成固定长度加强型密码（至少一个大写字母，一个小写字母，3个数字）

    Args:
        len (int): 密码长度

    Returns:
        str: 随机密码
    """
    assert (min_uppers+min_lowers+min_digits)<=len, f"总长度不能超过{len}"

    uppers = [secrets.choice(string.ascii_letters).upper() for i in range(min_uppers)]
    lowers = [secrets.choice(string.ascii_letters).lower() for i in range(min_lowers)]
    digits = [secrets.choice(string.digits) for i in range(min_digits)]
    others = [secrets.choice(alphanum) for i in range(len-min_uppers-min_lowers-min_digits)]
    passwd = uppers+lowers+digits+others
    random.shuffle(passwd)
    return ''.join(passwd)

    # while True:
    #     passwd = ''.join(secrets.choice(alphanum) for i in range(len))
    #     if (any(c.islower() for c in passwd)
    #             and any(c.isupper() for c in passwd)
    #             and len(c.isdigit() for c in passwd) >= 3):
    #         break

    # return passwd

def token(len: int) -> str:
    """生成固定长度的token

    Args:
        len (int): token长度

    Returns:
        str: 随机token
    """
    return secrets.token_urlsafe(len)

def md5(s: str, salt:bytes=b"x2trade") -> str:
    """生成md5摘要

    Args:
        s (str): 字符串
        salt(str): 字符串

    Returns:
        str: md码
    """
    h1 = hmac.new(salt, s.encode(), digestmod=hashlib.md5)
    return h1.hexdigest()

def encrypt(data, secret):
    """使用AES_CBC加密

    Args:
        data (bytes): 待加密字节串，可以encode为utf8
        secrect(bytes): 对称密钥

    Returns:
        bytes: 加密密文
    """
    iv = Random.new().read(AES.block_size)
    cipher = AES.new(secret, AES.MODE_CBC, iv)
    data = cipher.encrypt(pad(data, AES.block_size))
    data = iv + data
    return (data)
 
def decrypt(data, secret):
    """使用AES_CBC解密

    Args:
        data (bytes): 已加密字节串
        secrect(bytes): 对称密钥

    Returns:
        bytes: 解密密文
    """
    if len(data) <= AES.block_size:
        return (data)
    iv = data[:AES.block_size]
    cipher = AES.new(secret, AES.MODE_CBC, iv)
    data  = unpad(cipher.decrypt(data[AES.block_size:]),AES.block_size)
    return (data)    

def b64decode(data:bytes):
    """base64编码

    Args:
        data (bytes): 待编码字节串

    Returns:
        bytes: 经base64编码的ascii字节串
    """
    return base64.b64decode(data)

def b64encode(data:bytes):
    """base64解码

    Args:
        data (bytes): 经base64编码的ascii字节串

    Returns:
        bytes: 经base64解码的字节串
    """
    return base64.b64encode(data)